import json
import os
import argparse

import numpy as np

import DocTer.fuzzer.util as util


def get_array_str(arr):
    if len(arr) == 1:
        return str(arr[0])

    arr_str = "("
    for i in range(len(arr)):
        if i > 0:
            arr_str += ", "
        arr_str += str(arr[i])
    arr_str += ")"

    return arr_str


def generate_seed(test_case, numpy_arr_path):
    nd_arrays = np.load(numpy_arr_path)
    param_list = test_case.keys()
    for param in param_list:
        val = test_case[param]
        if type(val) is str and val in nd_arrays:
            test_case[param] = nd_arrays[val]
            # print(nd_arrays[val])
    return test_case


def multi_inputs_generator(test_case, target_config, model_path, om_path, index):
    all_name = target_config.get("title")
    constraints = target_config.get("constraints")
    input_params = ""
    func_params = ""
    # Is there an input existed
    flag = False
    # input dict
    input_dict = {
        "test_case_file_name": "",
        "inputs": []
    }
    # analyze test case
    param_list = test_case.keys()
    for param in param_list:
        seed = test_case[param]
        # ignore the name of input
        if param == "name":
            continue
        if type(seed) is int and seed > 10000:
            continue

        # np.array is considered as an input
        if type(seed) is np.ndarray:
            is_too_large = False
            space = 1
            for s in seed.shape:
                if space > int(1e5):
                    is_too_large = True
                    break
                else:
                    space *= s
            if is_too_large is True:
                continue
            input_dict["inputs"].append({param: param + "_" + str(index)})
            shape_str = get_array_str(seed.shape)
            type_str = "tf." + str(seed.dtype)
            input_params += "    %s = tf.compat.v1.placeholder(%s, shape=%s, name=\"%s\")\n" \
                            % (param, type_str, shape_str, param)
            func_params += "%s=%s, " % (param, param)
            flag = True
        else:
            if type(seed) == str:
                input_params += "    %s = \"%s\"\n" % (param, seed)
            else:
                seed = str(seed)
                input_params += "    %s = %s\n" % (param, seed)
            func_params += "%s=%s, " % (param, param)

    # Does API have attribute 'name'
    func_params = func_params[0:len(func_params) - 2]
    output_params = "    output_data = %s(%s, name=\"output\")\n" % (all_name, func_params)
    if "name" not in constraints:
        output_params = "    r = %s(%s)\n" % (all_name, func_params)
        output_params += "    output_data = tf.constant([r], name=\"output\")\n"

    # If there is no input data
    if flag is False:
        input_params += "    input = tf.compat.v1.placeholder(tf.int32, shape=1, name=\"input\")\n"

    # create a new onnx builder
    model_name = all_name + "_" + str(index)
    input_dict["test_case_file_name"] = model_name
    with open("/home/ubuntu/onnx_samples/onnx_transformer/tf_cann_generator.py", "w+") as f:
        f.write("import tensorflow as tf\n"
                "import os\n"
                "from tensorflow.python.framework.graph_util import convert_variables_to_constants\n\n"
                "model_root_path = '%s'\n\n\n" % model_path +
                "def create_model():\n"
                "    test_name = '%s'\n" % model_name +
                "    tf.compat.v1.disable_eager_execution()\n" +
                input_params + output_params +
                "    with tf.compat.v1.Session() as sess:\n"
                "        sess.run(tf.compat.v1.global_variables_initializer())\n"
                "        graph = convert_variables_to_constants(sess, sess.graph_def, [\"output\"])\n"
                "        tf.io.write_graph(graph, '.', os.path.join(model_root_path, test_name + '.pb'), as_text=False)\n"
                "        print('Create Model Successful.')\n"
                "        print('Path: ', os.path.join(model_root_path, test_name + '.pb'))\n\n"
                "    tf.compat.v1.reset_default_graph()\n\n\n"
                "if __name__ == \"__main__\":\n"
                "    create_model()")

    # execute the onnx builder
    full_model_path = os.path.join(model_path, model_name + ".pb")
    full_om_path = os.path.join(om_path, model_name)
    os.system("cd ~/onnx_samples/onnx_transformer/scripts;chmod +x tf2cann.sh;./tf2cann.sh %s %s;" %
             (full_model_path, full_om_path))

    return input_dict


def generate_models(seed_path, numpy_arr_path, model_path, om_path, input_path, config):
    test_inputs = []
    with open(seed_path, "r") as f:
        seeds = json.load(f)
    for i in range(len(seeds)):
        seed = seeds[i]
        if numpy_arr_path != "":
            seed = generate_seed(seed, numpy_arr_path)
        # print(case)
        input_dict = multi_inputs_generator(seed, config, model_path, om_path, i)
        test_inputs.append(input_dict)
    full_input_path = os.path.join(input_path, config.get("title") + "_inputs.json")
    # save input dicts to json
    with open(full_input_path, "w+") as f:
        json.dump(test_inputs, f)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--seed_path", "-s", default="/home/ubuntu/onnx_samples/onnx_transformer/source")
    parser.add_argument("--json_path", "-j", default="tf.math.abs_seeds.json")
    parser.add_argument("--npz_path", "-n", default="tf.math.abs_seeds.npz")
    parser.add_argument("--model_path", "-m", default="/home/ubuntu/onnx_samples/onnx_transformer/source")
    parser.add_argument("--om_path", "-o", default="/home/ubuntu/onnx_samples/onnx_transformer/source")
    parser.add_argument("--input_path", "-i", default="/home/ubuntu/onnx_samples/onnx_transformer/source")
    args = parser.parse_args()
    api_name = args.json_path[:-11]
    print("[TF Model] Api Name: " + api_name)
    config = util.read_yaml("/home/ubuntu/onnx_samples/onnx_transformer/DocTer/constraints/tf_raw_ops/%s.yaml" % api_name)
    generate_models(os.path.join(args.seed_path, args.json_path),
                    os.path.join(args.seed_path, args.npz_path),
                    args.model_path, args.om_path, args.input_path, config)
